import math
from collections import defaultdict
from typing import Dict, Tuple, Optional

from per.mnxc.survival_evolution.constants.constants import GRID_CELL_SIZE
from per.mnxc.survival_evolution.po.organism import Organism


class SpaceGridCell:
    def __init__(self, pos: Tuple[int, int]):
        self.pos = pos
        self.organisms: defaultdict[type, list[Organism]] = defaultdict(list[Organism])

class SpaceGrid:
    """空间网格加速结构"""

    def __init__(self, world_size: Tuple[int, int], cell_size: int = GRID_CELL_SIZE):
        self.cell_size = cell_size
        # 创建二维列表，用于存储网格单元格
        self.grid: Dict[Tuple[int, int], SpaceGridCell] = {}
        self._init_grid(world_size)
        # 记录每个生物所在的网格单元格
        self.organism_cell: Dict[Organism, SpaceGridCell] = {}
        # 预计算相邻单元格偏移量（包括自身单元格）
        self.neighbor_offsets = [(-1, -1), (0, -1), (1, -1),
                                 (-1, 0), (0, 0), (1, 0),
                                 (-1, 1), (0, 1), (1, 1)]
    def _init_grid(self, world_size: Tuple[int, int]):
        for x in range(0, world_size[0], self.cell_size):
            for y in range(0, world_size[1], self.cell_size):
                self.grid[(x, y)] = SpaceGridCell((x,y))

    def _get_cell_key(self, x: float, y: float) -> Tuple[int, int]:
        """将坐标转换为网格坐标"""
        return int(x // self.cell_size), int(y // self.cell_size)

    def refresh_organism(self, org: Organism):
        if org is None:
            return
        if org.is_alive is False:
            """从网格中删除生物"""
            cell = self.organism_cell.get(org)
            if cell is not None:
                cell.organisms[type(org)].remove(org)
                self.organism_cell.pop(org)
            return
        old_cell = self.organism_cell.get(org)
        key = self._get_cell_key(org.x, org.y)
        cell: SpaceGridCell = self.grid[key]
        if old_cell is None:
            """添加生物到网格"""
            cell.organisms[type(org)].append(org)
            self.organism_cell[org] = cell
        if old_cell is not cell:
            """更新生物位置"""
            old_cell.organisms[type(org)].remove(org)
            cell.organisms[type(org)].append(org)
            self.organism_cell[org] = cell
    def find_intersect_cells(self, org: Organism) -> list[SpaceGridCell]:
        cell = self.organism_cell.get(org)
        if cell is None:
            return []
        l_d = org.x % self.cell_size
        t_d = org.y % self.cell_size
        r_d = self.cell_size - l_d
        b_d = self.cell_size - t_d
        lt_d = math.hypot(l_d, t_d)
        lb_d = math.hypot(l_d, b_d)
        rt_d = math.hypot(r_d, b_d)
        rb_d = math.hypot(r_d, t_d)
        d_list = [lt_d,t_d,rt_d,l_d,0,r_d,lb_d,b_d,rb_d]
        for i, d in enumerate(d_list):
            if d < org.view:
                dx,dy = self.neighbor_offsets[i]
                c = self.grid[(cell.pos[0]+dx,cell.pos[1]+dy)]
                if c is not None:
                    yield c



    def find_nearest(self, org: Organism, target_type: type) -> Optional[Organism]:
        """查找最近的特定类型生物（优化版）"""
        closest = None
        min_distance = float('inf')
        # 检查当前单元格及相邻的8个单元格
        for candidates in self.find_intersect_cells(org):
            for candidate in candidates.organisms.get(target_type, []):
                if candidate is org or candidate.is_alive is False:
                    continue
                distance = math.hypot(candidate.x - org.x, candidate.y - org.y)
                if distance <= org.view and distance < min_distance:
                    closest = candidate
                    min_distance = distance
        return closest
